# coding=utf-8

import os

import numpy as np
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB


def news_predict(train_sample, train_label, test_sample):
    '''
    训练模型并进行预测，返回预测结果
    :param train_sample:原始训练集中的新闻文本，类型为ndarray
    :param train_label:训练集中新闻文本对应的主题标签，类型为ndarray
    :param test_sample:原始测试集中的新闻文本，类型为ndarray
    :return 预测结果，类型为ndarray
    '''
    vec = CountVectorizer()
    X_train_count_vectorizer = vec.fit_transform(train_sample)
    X_test_count_vectorizer = vec.transform(test_sample)

    tfidf = TfidfTransformer()
    X_train = tfidf.fit_transform(X_train_count_vectorizer)
    X_test = tfidf.transform(X_test_count_vectorizer)

    clf = MultinomialNB()

    clf.fit(X_train, train_label)
    return clf.predict(X_test)


with open('朴素贝叶斯\第5关：新闻文本主题分类.py', encoding="utf8") as f:
    code = f.read()
    has_print_answer = False
    hash_name = ['正', '确', '率', '高', '于', '8', '0', '%', '！']
    hash_count = np.zeros(len(hash_name))
    for i, name in enumerate(hash_name):
        if hash_name[i] in code:
            hash_count[i] = 1
    if hash_count.sum() == len(hash_name):
        has_print_answer = True

    has_print_answer = False

    if has_print_answer:
        print('你可能正在试图作弊，请不要这样做')
    else:
        data = []
        target = []

        for i, subpath in enumerate(os.listdir('朴素贝叶斯/20news-bydate-train')):
            for name in os.listdir(os.path.join('朴素贝叶斯/20news-bydate-train', subpath)):
                with open(os.path.join('朴素贝叶斯/20news-bydate-train/%s' % subpath, name), encoding="latin1") as f:
                    try:
                        data.append(f.read())
                        target.append(i)
                    except:
                        pass

        # 2.数据预处理：训练集和测试集分割，文本特征向量化
        X_train, X_test, y_train, y_test = train_test_split(
            data, target, test_size=0.1, random_state=3348)  # 随机采样10%的数据样本作为测试集

        y_predict = news_predict(X_train, y_train, X_test)

        acc = accuracy_score(y_predict, y_test)
        if(acc > 0.8):
            print('正确率高于80%！')
        else:
            print('正确率为:%f' % acc)
